import opensim as osim
import os

input_path = "das3_4x.osim"
output_path = "das3_with_left.osim"

# 读取原模型
model = osim.Model(input_path)
model.initSystem()

# 获取所有 body 名称
bodyset = model.getBodySet()
body_names = [body.getName() for body in bodyset]
right_bodies = [b for b in body_names if b not in ['ground', 'thorax']]

print("右臂刚体：", right_bodies)

# ---- 镜像刚体 ----
for bname in right_bodies:
    body_r = bodyset.get(bname)
    body_l = osim.Body(body_r)
    if bname.endwith("_r"):
        new_name = bname.replace("_r", "_l")
    elif bname[-1].isdigit() and bname[-1] == '1':
        new_name = bname[:-1] + '3'
    elif bname[-1].isdigit() and bname[-1] == '2':
        new_name = bname[:-1] +'4'
    body_l.setName(new_name)

    # 镜像质心位置
    mc = body_l.getMassCenter()
    body_l.setMassCenter(osim.Vec3(-mc.get(0), mc.get(1), mc.get(2)))

    # 镜像附加几何体
    geomset = body_l.get_AttachedGeometrySet()
    for i in range(geomset.getSize()):
        geom = geomset.get(i)
        geom.setName(geom.getName().replace("_r", "_l"))
        geom.set_scale_factors(osim.Vec3(-1, 1, 1))

    model.addBody(body_l)

# ---- 镜像关节 ----
jointset = model.getJointSet()
for i in range(jointset.getSize()):
    joint = jointset.get(i)
    name = joint.getName()
    if "_r" in name:
        j_new = osim.CustomJoint(joint)
        j_new.setName(name.replace("_r", "_l"))

        loc = j_new.getLocationInParent()
        j_new.setLocationInParent(osim.Vec3(-loc.get(0), loc.get(1), loc.get(2)))
        orient = j_new.getOrientationInParent()
        j_new.setOrientationInParent(osim.Vec3(orient.get(0), -orient.get(1), -orient.get(2)))

        model.addJoint(j_new)

# ---- 镜像肌肉 ----
forceset = model.getForceSet()
for i in range(forceset.getSize()):
    f = forceset.get(i)
    if not isinstance(f, osim.Muscle):
        continue

    path = f.getGeometryPath()
    npts = path.getPathPointSet().getSize()
    right_path = True
    for j in range(npts):
        p = path.getPathPointSet().get(j)
        body_name = p.getParentFrame().getName()
        if not body_name.endswith("_r"):
            right_path = False
            break

    if right_path:
        new_muscle = osim.Millard2012EquilibriumMuscle(f)
        name = f.getName()
        if "_r" in name:
            new_muscle.setName(name.replace("_r", "_l"))
        else:
            new_muscle.setName(name + "_l")

        # 镜像路径点
        new_path = new_muscle.getGeometryPath()
        for j in range(npts):
            p = path.getPathPointSet().get(j)
            loc = p.getLocation()
            mirrored_loc = osim.Vec3(-loc.get(0), loc.get(1), loc.get(2))
            new_path.updPathPointSet().get(j).setLocation(mirrored_loc)

        model.addForce(new_muscle)

# 输出模型
model.printToXML(output_path)
print("✅ 新模型已保存：", output_path)
